# 用矩阵块求解stokes方程
# Reference He's Book 5.3 稳态Navier-Stokes方程
# https://jsdokken.com/FEniCS-workshop/src/deep_dive/mixed_problems.html
# 和这个问题基本一样
from mpi4py import MPI
import numpy as np

import ufl
from basix.ufl import element, mixed_element
from dolfinx import default_real_type, fem, la, log, mesh
from dolfinx.fem import (
    Constant,
    Function,
    dirichletbc,
    form,
    functionspace,
    petsc
)
from petsc4py import PETSc
from dolfinx.fem.petsc import assemble_matrix_block, assemble_vector_block, set_bc, LinearProblem
from dolfinx.io import XDMFFile
from dolfinx.mesh import CellType, create_rectangle
from ufl import div, dx, inner, nabla_grad, dot, TestFunction, TrialFunction, TestFunctions, TrialFunctions, SpatialCoordinate
from dolfinx import la
import GetConvergenceRate as gc


# Solution
# Namespace: u_old表示前一次迭代结果 u为本次迭代结果


def getStiffnessMatrix(A: PETSc.Mat):
    """
    获取刚度矩阵的值(np.array)
    """
    # 需要导入
    # from petsc4py import PETSc
    m, n = A.getSize()
    matrix_data = np.zeros((m, n))

    # 获取矩阵的所有非零元素
    row, col, val = A.getValuesCSR()

    # 将CSR格式的数据转换为密集矩阵
    for i in range(m):
        for j in range(row[i], row[i + 1]):
            matrix_data[i, col[j]] = val[j]

    return matrix_data


def u_x(x):
    """
    x方向真实解
    """
    return 20 * x[0] ** 2 * (x[0] - 1) ** 2 * x[1] * (x[1] - 1) * (2 * x[1] - 1)


def u_y(x):
    """
    y方向真实解
    """
    return -20 * x[0] * (x[0] - 1) * (2 * x[0] - 1) * x[1] ** 2 * (x[1] - 1) ** 2


def pressure(x):
    """
    压强的真实解
    """
    return 10 * (2 * x[0] - 1) * (2 * x[1] - 1)


def true_velocity(x):
    """
    获取真实速度
    return: List
    """
    return [u_x(x), u_y(x)]


def f_x(x):
    return - 240.0 * nu * x[0] ** 4 * x[1] + 120.0 * nu * x[0] ** 4 + 480.0 * nu * x[0] ** 3 * x[1] - 240.0 * nu * x[0] ** 3 - 480.0 * nu * x[0] ** 2 * x[1] ** 3 + 720.0 * nu * x[0] ** 2 * x[1] ** 2 - 480 * nu * x[0] ** 2 * \
        x[1] + 120.0 * nu * x[0] ** 2 + 480.0 * nu * x[0] * x[1] ** 3 - 720.0 * nu * x[0] * x[1] ** 2 + 240.0 * \
        nu * x[0] * x[1] - 80.0 * nu * x[1] ** 3 + 120.0 * \
        nu * x[1] ** 2 - 40.0 * nu * x[1] + 40 * x[1] - 20


def f_y(x_list):
    return 480.0 * nu * x[0] ** 3 * x[1] ** 2 - 480.0 * nu * x[0] ** 3 * x[1] + 80.0 * nu * x[0] ** 3 - 720.0 * nu * x[0] ** 2 * x[1] ** 2 + 720.0 * nu * x[0] ** 2 * x[1] - 120.0 * nu * x[0] ** 2 + 240.0 * nu * x[0] * \
        x[1] ** 4 - 480.0 * nu * x[0] * x[1] ** 3 + 480 * nu * x[0] * x[1] ** 2 - 240.0 * nu * x[0] * x[1] + 40.0 * \
        nu * x[0] - 120.0 * nu * x[1] ** 4 + 240.0 * nu * \
        x[1] ** 3 - 120.0 * nu * x[1] ** 2 + 40 * x[0] - 20


def f(x):
    return [f_x(x), f_y(x)]


def D(u):
    return 0.5 * (ufl.nabla_grad(u) + ufl.nabla_grad(u).T)


def get_norm(sol: PETSc.Vec, print_=True, save_=True):
    """计算范数

    Args:
        sol (PETSc.Vec): 基于分别定义空间完成求解的解向量
        print_ (bool, optional): 是否打印输出
        save_ (bool, optional): 是否保存结果
    """
    # 构造真解
    u_true = Function(V)
    p_true = Function(Q)
    u_true.interpolate(true_velocity)
    p_true.interpolate(pressure)

    # H1范数 按H1半范进行计算
    u_error = Function(V)
    # 利用offset分解向量
    offset = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
    u_h = sol[:offset]
    p_h = sol[offset:]
    u_error.x.array[:] = u_h - u_true.x.array
    p_error = Function(Q)
    p_error.x.array[:] = p_h - p_true.x.array

    # Compute norms of the u and p vectors
    u_comm = u_true.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    u_h1_error = form((ufl.grad(u_error)**2) * ufl.dx)
    u_l2_error = form(u_error**2 * ufl.dx)
    u_h1_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_h1_error), MPI.SUM))
    u_l2_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_l2_error), MPI.SUM))  # 并行求和 返回float

    p_comm = p_true.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    p_h1_error = form((ufl.grad(p_error)**2) * ufl.dx)
    p_l2_error = form(p_error**2 * ufl.dx)
    p_h1_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_h1_error), MPI.SUM))
    p_l2_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_l2_error), MPI.SUM))  # 并行求和 返回float

    # Linf 比较节点值 离散后区域无法计算理论最大值
    u_linf_E, p_linf_E = la.norm(
        u_error.x, la.Norm.linf), la.norm(p_error.x, la.Norm.linf)
    data_ = [
        n_list[idx],
        u_h1_E, p_h1_E,
        u_l2_E, p_l2_E,
        u_linf_E, p_linf_E,
    ]

    if MPI.COMM_WORLD.rank == 0:
        if save_:  # 是否写入为excel
            gc.save_data(FILENAME, data_)
        if print_:
            print(
                f"H1 norm for velocity: {u_h1_E}, pressure: {p_h1_E}.")
            print(
                f"L2 norm for velocity: {u_l2_E}, pressure: {p_l2_E}.")
            print(
                f"Linf norm for velocity: {u_linf_E}, pressure: {p_linf_E}.")


def solve_linear_system(
    A: PETSc.Mat,                # 分块矩阵
    b: PETSc.Vec,                # 右端项向量
    msh: mesh.Mesh,              # 网格对象
    ksp_type: str = "preonly",      # 求解器类型 (gmres, minres, bicgstab 等)
    pc_type: str = "lu",  # 预条件器类型
) -> None:
    """
    使用 KSP 求解器求解分块线性系统 Ax = b 并更新混合函数解。

    参数:
        A: 分块矩阵 (assembled with assemble_matrix_block)
        b: 右端项向量 (assembled with assemble_vector_block)
        msh: 网格对象
        ksp_type: KSP 迭代方法类型
        pc_type: 预条件器类型
        rtol: 相对残差容限
        max_iter: 最大迭代次数
    """
    # 创建并配置 KSP 求解器
    ksp = PETSc.KSP().create(msh.comm)
    ksp.setOperators(A)  # 将系数矩阵 A 设置为 KSP 求解器的 主矩阵（Operator）
    ksp.setType(ksp_type)

    # 配置预条件器
    pc = ksp.getPC()
    pc.setType(pc_type)
    sys = PETSc.Sys()
    use_superlu = PETSc.IntType == np.int64
    if sys.hasExternalPackage("mumps") and not use_superlu:
        pc.setFactorSolverType("mumps")
        pc.setFactorSetUpSolverType()
        # icntl=24 允许处理奇异矩阵 =25禁用错误检测
        pc.getFactorMatrix().setMumpsIcntl(icntl=24, ival=1)
        pc.getFactorMatrix().setMumpsIcntl(icntl=25, ival=0)
    else:
        pc.setFactorSolverType("superlu_dist")
    x = A.createVecLeft()

    # stiff_mat = getStiffnessMatrix(A=A)  # 刚度矩阵太大会oom(n=128)
    # print(np.linalg.cond(stiff_mat))
    load_vec = b.getArray()  # 载荷向量可以直接获取
    for (i, e) in enumerate(load_vec):
        if e > 1e-10 or e < -1e-10:
            print(f"Warning: Large value in load vector: {e} idx = {i}")
    # 零空间
    # null_vec = A.createVecLeft()
    # offset = Q.dofmap.index_map.size_local * Q.dofmap.index_map_bs
    # # print(f"{V.dofmap.index_map.size_local}, {V.dofmap.index_map_bs}")
    # null_vec.array[offset:] = 1.0
    # null_vec.normalize()
    # nsp = PETSc.NullSpace().create(vectors=[null_vec])
    # assert nsp.test(A)
    # A.setNullSpace(nsp)

    # 求解线性系统
    try:
        ksp.solve(b, x)
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT,
                      mode=PETSc.ScatterMode.FORWARD)
        print(
            f"KSP Solver converged in {ksp.getIterationNumber()} iterations.")
    except PETSc.Error as e:
        print(f"KSP Solver failed: {e}")
        raise

    # 释放资源
    ksp.destroy()
    return x


FILENAME = f"ns-block.xlsx"
gc.create_file(FILENAME)
# Params
n_list = [[16, 16],
          [32, 32],
          [64, 64],
          [128, 128],
          ]  # 剖分列表
idx = int(input("请输入网格划分序号："))  # 0-4
nu = 1
rtol = 1e-8
max_iter = 1000

msh = create_rectangle(
    MPI.COMM_WORLD, [np.array([0, 0]), np.array(
        [1, 1])], n=n_list[idx], cell_type=CellType.triangle
)

# Function Space
P2 = element(
    "Lagrange", msh.basix_cell(), degree=2, shape=(msh.geometry.dim,), dtype=default_real_type
)
P1 = element("Lagrange", msh.basix_cell(), degree=1, dtype=default_real_type)
V, Q = functionspace(msh, P2), functionspace(msh, P1)  # 不使用混合元

x = SpatialCoordinate(msh)  # 定义空间坐标

u_bc = Function(V)  # 真实速度函数空间
p_bc = Function(Q)  # 真实压力函数空间
u_bc.interpolate(true_velocity)
p_bc.interpolate(pressure)


# type u, p =  <class 'ufl.tensors.ListTensor'>, <class 'ufl.indexed.Indexed'>
# mixed_space使用TrialFunctions/TestFunctions 注意这里多了s
(u, p) = ufl.TrialFunction(V), ufl.TrialFunction(Q)
(v, q) = ufl.TestFunction(V), ufl.TestFunction(Q)


# BCs
fdim = msh.topology.dim - 1
facets = mesh.locate_entities_boundary(
    msh, fdim, lambda x: np.isclose(x[0], 1.0) | np.isclose(
        x[0], 0.0) | np.isclose(x[1], 1.0) | np.isclose(x[1], 0.0)
)
# 自由度的空间使用(V0, Q) 边界条件设置为V0
# 这里写法是固定的要记住!!!
v_dofs = fem.locate_dofs_topological(V, fdim, facets)
# 当传入value是插值函数 会自动导入函数空间 不需要设置空间
bc_vel = dirichletbc(u_bc, v_dofs)  # 速度边界条件

# 固定压力点
def marker(x): return np.isclose(x[0], 0.0) & np.isclose(x[1], 0)


p_dofs = fem.locate_dofs_geometrical(Q, marker)
bc_pressure = dirichletbc(p_bc, p_dofs)


# Weak Formulation
f_expr = ufl.as_vector([f_x(x), f_y(x)])  # 调整f为ufl格式
L = inner(f_expr, v) * dx  # 右端项

# LinearSolver求解
bcs = [bc_vel, bc_pressure]


# 矩阵分块逻辑 不要写错顺序！ f1 f2 f3 f4为对应的内积形式映射
# 如f1(u, v) = 2 * nu * inner(D(u), D(v)) * dx
#       u     p
# v   f1(u,v) f2(p, v)
# q   f3(u, q) f4(p, q)
a_uu = 2 * nu * inner(D(u), D(v)) * dx
a_vp = -inner(div(v), p) * dx
a_uq = -inner(div(u), q) * dx
a_pp = Constant(msh, 0.0) * inner(p, q) * dx  # 这里的Constant(msh, 0)的作用是指定积分区域
D
L_u = inner(f_expr, v) * dx  # 右端项
L_p = inner(Constant(msh, 0.0), q) * dx

a = form([[a_uu, a_vp], [a_uq, a_pp]])
L = form([L_u, L_p])
# assemble_matrix_block 和 assemble_vector_block函数自动处理边界条件
# 不使用混合元空间 一切正常 目前的观点是使用混合元空间+分块求解会导致一些问题

A = assemble_matrix_block(a, bcs=bcs)  # petsc4py.PETSc.Mat
L = assemble_vector_block(L, a, bcs=bcs)  # petsc4py.PETSc.Vec
A.assemble()
L.assemble()

# 求解 返回一个PETSc.Vec对象
sol = solve_linear_system(
    A=A,
    b=L,
    msh=msh,
    ksp_type="preonly",  # 可动态调整为 "minres" 等
)
# 计算误差
# get_norm(sol)
